#include <src/pcps.h>
#include <src/input.h>
#include <src/wavefunction.h>
#include <src/operator.h>
#include <src/cluster.h>

//--------------------------------------------------------------------------------
// pcps::Cluster::add_site -- adds a site to the cluster
//--------------------------------------------------------------------------------

void pcps::Cluster::add_site(const int site, const int offset) {

    // if the site is already present, raise an error
    if (_site_lookup.at(site) >= 0)
      throw pcps::Exception("cannot add site to Cluster object:  the site is already present");

    // record the site
    _sites.push_back(site);

    // record the offset
    _offsets.push_back(offset);

    // set the number of allowed occupations
    _n_allowed_occs.push_back(_configs_per_site);

    // set the number of the current occupation
    _occ_nums.push_back(0);

    // prepare a vector of allowed occupations
    std::vector<int> temp_vec(_configs_per_site);
    for (int i = 0; i < _configs_per_site; i++)
      temp_vec.at(i) = i;
    _allowed_occs.push_back(temp_vec);

    // record the site in the lookup table
    _site_lookup.at(site) = _sites.size() - 1;

    // record the site in the ket occupation vector
    _ket_occs.at(site) = _allowed_occs.at(_site_lookup.at(site)).at(_occ_nums.at(_site_lookup.at(site)));

    // count the site
    _size++;

}

//--------------------------------------------------------------------------------
// pcps::Cluster::Cluster -- Constructs a Cluster object of a given number of sites
//--------------------------------------------------------------------------------

pcps::Cluster::Cluster(const int n_lattice_sites, const int configs_per_site)
  : _configs_per_site(configs_per_site),
    _size(0),
    _site_lookup(n_lattice_sites, -1),
    _ket_occs(n_lattice_sites, -1)
{
  // add each site with a zero offset
  for (int s = 0; s < n_lattice_sites; s++)
    this->add_site(s, 0);

  // initialize the occupation vector
  this->reset_occs();
}

//--------------------------------------------------------------------------------
// pcps::Cluster::Cluster -- Constructs a Cluster object from a set of sites
//--------------------------------------------------------------------------------

pcps::Cluster::Cluster(const int n_lattice_sites, const int configs_per_site, const std::set<int> & sites)
  : _configs_per_site(configs_per_site),
    _size(0),
    _site_lookup(n_lattice_sites, -1),
    _ket_occs(n_lattice_sites, -1)
{
  // add each site with a zero offset
  for (std::set<int>::const_iterator site = sites.begin(); site != sites.end(); site++)
    this->add_site(*site, 0);

  // initialize the occupation vector
  this->reset_occs();
}

//--------------------------------------------------------------------------------
// pcps::Cluster::Cluster -- Constructs a Cluster object from a vector of sites
//--------------------------------------------------------------------------------

pcps::Cluster::Cluster(const int n_lattice_sites, const int configs_per_site, const std::vector<int> & sites)
  : _configs_per_site(configs_per_site),
    _size(0),
    _site_lookup(n_lattice_sites, -1),
    _ket_occs(n_lattice_sites, -1)
{
  // add each site with a zero offset
  for (std::vector<int>::const_iterator site = sites.begin(); site != sites.end(); site++)
    this->add_site(*site, 0);

  // initialize the occupation vector
  this->reset_occs();
}

//--------------------------------------------------------------------------------
// pcps::Cluster::Cluster -- Constructs a Cluster object containing the sites
//                           around the operator.
//                           The sites included are:
//                            (1) The sites touched by the operator.
//                            (2) The sites whose occupations are fixed, as
//                                described by fixed_bra_occs and fixed_ket_occs.
//                            (3) Any site that shares a correlator with an
//                                operator site that has a non-zero offset.
//--------------------------------------------------------------------------------

template <class S>
pcps::Cluster::Cluster(const pcps::Input & userinp,
                       const pcps::Wavefunction<S> & wfn,
                       const typename pcps::OpBase<S>::const_ptr & op,
                       const pcps::Cluster & fixed_bra_occs,
                       const pcps::Cluster & fixed_ket_occs,
                       const std::vector<int> & pert_corr_pos)
  : _configs_per_site(userinp.configs_per_site()),
    _size(0),
    _site_lookup(userinp.nsites(), -1),
    _ket_occs(userinp.nsites(), -1)
{

  // Determine the particle types
  const bool particles_are_spins = ( userinp.particle_type() == pcps::keyword::Part_Spin );
  const bool particles_are_fermions = (    userinp.particle_type() == pcps::keyword::Part_Fermion
                                        || userinp.particle_type() == pcps::keyword::Part_Spinless_Fermion );

  // Include sites touched by the operator
  for (int i = 0; i < op->sites().size(); i++) {

    // get the site and offset
    const int site = op->sites().at(i);
    const int offset = op->offsets().at(i);

    // check that the site has not already been added
    if (_site_lookup.at(site) >= 0)
      throw pcps::Exception("operators should not have repeated site indeces");

    this->add_site(site, offset);

  }

  // Indclude sites whose ket occupations are fixed
  for (int i = 0; i < fixed_ket_occs.sites().size(); i++) {

    // get the site and occupation
    const int site = fixed_ket_occs.sites().at(i);
    const int occ = fixed_ket_occs.ket_occ(site);

    // if the site is new, add it to the cluster with a zero offset
    if (_site_lookup.at(site) < 0)
      this->add_site(site, 0);

    // Fix the site's ket occupation by deleting ket occupations that are not allowed
    const int site_pos = _site_lookup.at(site);
    for (int i = 0; i < _n_allowed_occs[site_pos]; )
      if (_allowed_occs[site_pos][i] != occ) {
        for (int j = i; j < _configs_per_site - 1; j++)
          _allowed_occs[site_pos][j] = _allowed_occs[site_pos][j+1];
        _n_allowed_occs[site_pos]--;
      } else {
        i++;
      }

  }

  // Indclude sites whose bra occupations are fixed
  for (int i = 0; i < fixed_bra_occs.sites().size(); i++) {

    // get the site and occupation
    const int site = fixed_bra_occs.sites().at(i);
    const int occ = fixed_bra_occs.ket_occ(site);

    // if the site is new, add it to the cluster with a zero offset
    if (_site_lookup.at(site) < 0)
      this->add_site(site, 0);

    // Fix the site's bra occupation by deleting ket occupations that are not allowed
    const int site_pos = _site_lookup.at(site);
    for (int i = 0; i < _n_allowed_occs[site_pos]; )
      if (_allowed_occs[site_pos][i] + _offsets[site_pos] != occ) {
        for (int j = i; j < _configs_per_site - 1; j++)
          _allowed_occs[site_pos][j] = _allowed_occs[site_pos][j+1];
        _n_allowed_occs[site_pos]--;
      } else {
        i++;
      }

  }

  // For each of the operator's sites, eliminate occupations for which the action value is always zero.
  // Note that if the number of cluster configurations becomes zero, there is no need to continue with this elimination.
  for (int i = 0; i < op->sites().size() && this->nconfigs() > 0; i++) {

    // get the site and its position
    const int site = op->sites().at(i);
    const int site_pos = _site_lookup[site];

    // iterate over the occupations of the site in question
    for (int site_occ_num = 0; site_occ_num < _n_allowed_occs[site_pos]; ) {

      // prepare vectors for the operator's sites' ket occupations
      std::vector<int> occs(op->sites().size());
      std::vector<int> occ_nums(op->sites().size());
      for (int j = 0; j < op->sites().size(); j++) {
        if (j == i)
          occ_nums.at(j) = site_occ_num;
        else
          occ_nums.at(j) = 0;
      }

      // initialize a flag to tell if any of the action values are not zero
      bool all_zero = true;

      // find the action value for each configuration of the other sites
      while (true) {

        // compute the operator's sites' occupations
        for (int j = 0; j < op->sites().size(); j++)
          occs[j] = _allowed_occs[ _site_lookup[op->sites()[j]] ][occ_nums[j]];

        // check if the action value is nonzero
        if (op->action_value(particles_are_spins, particles_are_fermions, _configs_per_site, &occs[0]) != 0.0) {
          all_zero = false;
          break;
        }

        // increment to the next occupation of the other sites
        int q = op->sites().size() - 1;
        while (q >= 0) {
          if (q == i)
            q--;
          else if ( ++occ_nums[q] == _n_allowed_occs[_site_lookup[op->sites()[q]]] ) {
            occ_nums[q] = 0;
            q--;
          } else
            break;
        }

        // stop iterating if all configurations have been processed
        if (q < 0) break;

      }

      // if all the action values are zero, remove the occupation
      if (all_zero) {
        for (int j = site_occ_num; j < _configs_per_site-1; j++)
          _allowed_occs[site_pos][j] = _allowed_occs[site_pos][j+1];
        _n_allowed_occs[site_pos]--;
      } else
        site_occ_num++;

    }

  }

  // Include sites that share a correlator with any of the offset sites
  const int old_size = _size;
  for (int i = 0; i < old_size; i++) {
    
    // get the site and offset
    const int site = _sites.at(i);
    const int offset = _offsets.at(i);

    // if the offset is nonzero, add all sites that share a correlator with the site
    if (offset != 0) {

      //// get a list of correlators touching the site
      //const std::vector<pcps::Correlator> & corr_vec = wfn.correlators_with_site(site);

      // for each correlator, add its sites to the cluster
      //for (std::vector<pcps::Correlator>::const_iterator corr = corr_vec.begin(); corr != corr_vec.end(); corr++)
      for (typename std::vector< pcps::Correlator<S> >::const_iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++)
        if (std::count(corr->site_vec().begin(), corr->site_vec().end(), site) > 0)
          for (std::vector<int>::const_iterator corr_site = corr->site_vec().begin(); corr_site != corr->site_vec().end(); corr_site++)
            if (_site_lookup.at(*corr_site) < 0)
              this->add_site(*corr_site, 0);

    }

  }

  // include sites in the perturbative correlators
  for (std::vector<int>::const_iterator pos = pert_corr_pos.begin(); pos != pert_corr_pos.end(); pos++)
    for (std::vector<int>::const_iterator site = wfn.pert_corrs().at(*pos).site_vec().begin(); site != wfn.pert_corrs().at(*pos).site_vec().end(); site++)
      if (_site_lookup.at(*site) < 0)
        this->add_site(*site, 0);

  // Initialize the occupation vector
  this->reset_occs();

}

// explicitly instantiate the template
template pcps::Cluster::Cluster
                         (const pcps::Input & userinp,
                          const pcps::Wavefunction<double> & wfn,
                          const pcps::OpBase<double>::const_ptr & op,
                          const pcps::Cluster & fixed_bra_occs,
                          const pcps::Cluster & fixed_ket_occs,
                          const std::vector<int> & pert_corr_pos);
template pcps::Cluster::Cluster
                         (const pcps::Input & userinp,
                          const pcps::Wavefunction< std::complex<double> > & wfn,
                          const pcps::OpBase< std::complex<double> >::const_ptr & op,
                          const pcps::Cluster & fixed_bra_occs,
                          const pcps::Cluster & fixed_ket_occs,
                          const std::vector<int> & pert_corr_pos);

//--------------------------------------------------------------------------------
// pcps::Cluster::split -- Splits a cluster into smaller clusters by finding the
//                         first site with multiple occupations and creating a
//                         new Cluster object for each occupation of that site.
//                         The splitting is continued recursively until all
//                         resulting Cluster objects have fewer than the specified
//                         number of configurations.
//--------------------------------------------------------------------------------

void pcps::Cluster::split(const int max_configs, std::list<pcps::Cluster> & subclusters) const {

  // initialize the list of smaller clusters
  subclusters.clear();

  // if the current cluster is small enough, then it does not need to be split
  if (this->nconfigs() <= max_configs) {
    subclusters.push_back(*this);
    return;
  }

  // otherwise, split up the cluster by the first site with multiple allowed occupations
  for (int i = 0; i < _size; i++) {
    if (_n_allowed_occs[i] > 1) {
      for (int occ_num = 0; occ_num < _n_allowed_occs[i]; occ_num++) {
        subclusters.push_back(*this);
        subclusters.rbegin()->_allowed_occs[i][0] = _allowed_occs[i][occ_num];
        //for (int j = 1; j < _configs_per_site; j++)
        //  subclusters.rbegin()->_allowed_occs[i][j] = -1;
        subclusters.rbegin()->_n_allowed_occs[i] = 1;
        subclusters.rbegin()->reset_occs();
      }
      break;
    }
  }

  // for any new cluster that is still not small enough, recursively split it into smaller clusters
  for (std::list<pcps::Cluster>::iterator clust = subclusters.begin(); clust != subclusters.end(); ) {
    if (clust->nconfigs() > max_configs) {
      std::list<pcps::Cluster> new_subcluster;
      clust->split(max_configs, new_subcluster);
      subclusters.erase(clust);
      for (std::list<pcps::Cluster>::iterator new_clust = new_subcluster.begin(); new_clust != new_subcluster.end(); new_clust++)
        subclusters.push_back(*new_clust);
      clust = subclusters.begin();
    }
    else
      clust++;
  }

}

//--------------------------------------------------------------------------------
// pcps::Cluster::send -- Send the cluster to another processor
//--------------------------------------------------------------------------------

void pcps::Cluster::send(const MPI::Comm & comm, const int dest, const int tag) const {

  comm.Send((const void *)(&_configs_per_site), 1, MPI::INT, dest, tag);
  comm.Send((const void *)(&_size), 1, MPI::INT, dest, tag);
  comm.Send((const void *)(&_finished), 1, MPI::BOOL, dest, tag);

  pcps::send_vec(comm, _site_lookup, dest, tag);
  pcps::send_vec(comm, _sites, dest, tag);
  pcps::send_vec(comm, _offsets, dest, tag);
  pcps::send_vec(comm, _n_allowed_occs, dest, tag);
  pcps::send_vec(comm, _occ_nums, dest, tag);
  pcps::send_vec(comm, _ket_occs, dest, tag);

  for (int i = 0; i < _size; i++)
    pcps::send_vec(comm, _allowed_occs.at(i), dest, tag);

}

//--------------------------------------------------------------------------------
// pcps::Cluster::recv -- Receive the cluster from another processor
//--------------------------------------------------------------------------------

void pcps::Cluster::recv(const MPI::Comm & comm, const int source, const int tag) {

  comm.Recv((void *)(&_configs_per_site), 1, MPI::INT, source, tag);
  comm.Recv((void *)(&_size), 1, MPI::INT, source, tag);
  comm.Recv((void *)(&_finished), 1, MPI::BOOL, source, tag);

  pcps::recv_vec(comm, _site_lookup, source, tag);
  pcps::recv_vec(comm, _sites, source, tag);
  pcps::recv_vec(comm, _offsets, source, tag);
  pcps::recv_vec(comm, _n_allowed_occs, source, tag);
  pcps::recv_vec(comm, _occ_nums, source, tag);
  pcps::recv_vec(comm, _ket_occs, source, tag);

  _allowed_occs.clear();
  _allowed_occs.resize(_size);

  for (int i = 0; i < _size; i++)
    pcps::recv_vec(comm, _allowed_occs.at(i), source, tag);

}
